package server

import (
	"fmt"
	"github.com/k0kubun/go-ansi"
	"github.com/schollz/progressbar/v3"
	"io"
	"log"
	"net/http"
	"os"
	"path"
	"strings"
	"sync"
)

type Downloader struct {
	concurrency int
	resume      bool
	bar         *progressbar.ProgressBar
}

func NewDownloader(concurrency int, resume bool) *Downloader {
	return &Downloader{concurrency: concurrency, resume: resume}
}

func (d *Downloader) Dowmload(strUrl, filename string) error {
	if filename == "" {
		filename = path.Base(strUrl)
	}
	head, err := http.Head(strUrl)
	if err != nil {
		return err
	}

	if head.StatusCode == http.StatusOK && head.Header.Get("Accept-Rangs") == "bytes" {
		return d.multiDownload(strUrl, filename, int(head.ContentLength))
	}

	return d.singleDownload(strUrl, filename)
}

func (d *Downloader) multiDownload(url string, filename string, contentlength int) error {
	//设置进度条
	d.SetBar(contentlength)

	parSize := contentlength / d.concurrency

	// 创建部分文件的存放目录
	parDir := d.getParDir(filename)
	os.Mkdir(parDir, 0777)
	defer os.RemoveAll(parDir)

	var wg sync.WaitGroup
	wg.Add(d.concurrency)

	rangerStart := 0

	for i := 0; i < d.concurrency; i++ {
		go func(i, rangerStart int) {
			defer wg.Done()
			rangerEnd := rangerStart + parSize
			// 最后一部分，总长度不能超过 ContentLength
			if i == d.concurrency-1 {
				rangerEnd = contentlength
			}
			download := 0

			if d.resume {
				partFileName := d.getParDirFilename(filename, i)
				content, err := os.ReadFile(partFileName)
				if err == nil {
					download = len(content)
				}
				d.bar.Add(download)
			}
			d.dowloadParial(url, filename, rangerStart+download, rangerEnd, i)
		}(i, rangerStart)
		rangerStart += parSize + 1
	}
	wg.Wait()
	d.merge(filename)
	return nil
}

func (d *Downloader) singleDownload(url string, filename string) error {
	resp, err := http.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	d.SetBar(int(resp.ContentLength))

	file, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	file.Close()
	buf := make([]byte, 32*1024)
	_, err = io.CopyBuffer(io.MultiWriter(file, d.bar), resp.Body, buf)
	return err
}

func (d *Downloader) SetBar(length int) {
	//填充进度条参数
	progressbar.NewOptions(
		length,
		progressbar.OptionSetWriter(ansi.NewAnsiStdout()),
		progressbar.OptionEnableColorCodes(true),
		progressbar.OptionShowBytes(true),
		progressbar.OptionSetWidth(50),
		progressbar.OptionSetDescription("下载中 . . . "),
		progressbar.OptionSetTheme(progressbar.Theme{
			Saucer:        "[green]=[reset]",
			SaucerHead:    "[green]>[reset]",
			SaucerPadding: " ",
			BarStart:      "[",
			BarEnd:        "]",
		}),
	)
}

func (d *Downloader) getParDir(filename string) string {
	return strings.SplitN(filename, ".", 2)[0]
}

// 构造部分文件名
func (d *Downloader) getParDirFilename(filename string, partNum int) string {
	parDir := d.getParDir(filename)
	return fmt.Sprintf("%s/%s-%d", parDir, filename, partNum)
}

func (d *Downloader) dowloadParial(url string, filename string, rangeStart int, rangeEnd int, i int) {
	if rangeStart >= rangeEnd {
		return
	}
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		log.Fatal(err)
	}
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", rangeStart, rangeEnd))
	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		log.Fatal(err)
	}
	defer resp.Body.Close()
	flags := os.O_CREATE | os.O_WRONLY
	if d.resume {
		flags |= os.O_APPEND
	}

	partFile, err := os.OpenFile(d.getParDirFilename(filename, i), flags, 0666)
	if err != nil {
		log.Fatal(err)
	}
	defer partFile.Close()

	buf := make([]byte, 32*1024)
	_, err = io.CopyBuffer(io.MultiWriter(partFile, d.bar), resp.Body, buf)
	if err != nil {
		if err == io.EOF {
			return
		}
		log.Fatal(err)
	}

}

func (d *Downloader) merge(filename string) error {
	destFile, err := os.OpenFile(filename, os.O_CREATE|os.O_WRONLY, 0666)
	if err != nil {
		return err
	}
	defer destFile.Close()

	for i := 0; i < d.concurrency; i++ {
		parFilename := d.getParDirFilename(filename, i)
		parFile, err := os.Open(parFilename)
		if err != nil {
			return err
		}
		io.Copy(destFile, parFile)
		parFile.Close()
		os.Remove(parFilename)
	}
	return nil
}
